# 堆相关的算法


class Heap:
    # 大根堆
    def __init__(self):
        self.arr = []
        self.heap_size = 0

    def add(self, item):
        """heap_insert从当前位置由下至上整理heap"""
        self.arr.append(item)
        self.heap_size += 1
        cur_index = self.heap_size - 1
        # f_index = (cur_index - 1) // 2
        # while f_index >= 0:
        #     if item > self.arr[f_index]:
        #         self._swap(f_index, cur_index)
        #         cur_index = f_index
        #         f_index = (cur_index - 1) // 2
        #     else:
        #         break
        while self.arr[cur_index] > self.arr[(cur_index - 1) // 2]:
            self._swap(cur_index, (cur_index - 1) // 2)
            cur_index = (cur_index - 1) // 2
            if cur_index == 0:
                break
        return self

    def heap_insert(self, index):
        """自下而上整理成大根堆"""
        while self.arr[index] > self.arr[(index - 1) // 2]:
            self._swap(index, (index - 1) // 2)
            index = (index - 1) // 2
            if index == 0:
                return

    def heapify(self, cur_index=0):
        """ 从当前index由上而下整理成heap"""
        last_index = self.heap_size - 1
        while cur_index * 2 + 1 < last_index or cur_index * 2 + 2 < last_index:
            if self.arr[cur_index * 2 + 1] > self.arr[cur_index] \
                    and self.arr[cur_index * 2 + 1] >= self.arr[cur_index * 2 + 2]:
                # print(self.arr[cur_index * 2 + 1])
                self._swap(cur_index, cur_index * 2 + 1)
                cur_index = cur_index * 2 + 1
            elif self.arr[cur_index * 2 + 2] > self.arr[cur_index] \
                    and self.arr[cur_index * 2 + 2] > self.arr[cur_index * 2 + 1]:
                # print(self.arr[cur_index * 2 + 2])
                self._swap(cur_index, cur_index * 2 + 2)
                cur_index = cur_index * 2 + 2
            else:
                return

    def pop_max(self):
        self._swap(0, self.heap_size - 1)
        self.heap_size -= 1
        self.heapify()
        return self.arr[self.heap_size]

    def _swap(self, i, j):
        self.arr[i], self.arr[j] = self.arr[j], self.arr[i]

    def __str__(self):
        return str(self.arr[:self.heap_size])


if __name__ == '__main__':
    heap = Heap()
    heap.add(0).add(2).add(5).add(1).add(4)
    heap.add(9)
    # heap.add(1).add(9).add(0)
    print(heap)
    print('pop max: ', heap.pop_max())
    print(heap)
    print('pop max: ', heap.pop_max())
    print(heap)
    print('pop max: ', heap.pop_max())
    print(heap)
    # heap.heapify()
